#!/usr/bin/env python3
################################################################
#
# makeapot:
#   create an analytic potential file for potfit
#
################################################################
#
#   Copyright 2013-2017 -  the potfit development team
#
#   https://www.potfit.net
#
#################################################################
#
#   This file is part of potfit.
#
#   potfit is free software; you can redistribute it and/or modify
#   it under the terms of the GNU General Public License as published by
#   the Free Software Foundation; either version 2 of the License, or
#   (at your option) any later version.
#
#   potfit is distributed in the hope that it will be useful,
#   but WITHOUT ANY WARRANTY; without even the implied warranty of
#   MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
#   GNU General Public License for more details.
#
#   You should have received a copy of the GNU General Public License
#   along with potfit; if not, see <http://www.gnu.org/licenses/>.
#
#################################################################

import argparse
import functions
import inspect
import random
import sys
import textwrap

INTERACTIONS = [
    ['pair', 'pair potentials', lambda x: paircol(x)],
    ['eam', 'embedded atom method', lambda x : paircol(x) + 2 * x],
    ['meam', 'modified embedded atom method', lambda x : 2 * paircol(x) + 3 * x],
    ['adp', 'angular dependent potential', lambda x : 3 * paircol(x) + 2 * x],
    ['stiweb', 'Stillinger-Weber potential', lambda x : 2 * paircol(x) + 1],
    ['tersoff', 'Tersoff potential', lambda x : x * x]
]


def paircol(n):
    return int(n * (n + 1) / 2)


def write_list():
    write_interaction_list()
    print("")
    write_pot_list()


def write_interaction_list():
    print('List of available interaction models:\n')
    length = 0
    for item in INTERACTIONS:
        length = max(length, len(item[0]))
    for item in INTERACTIONS:
        print("  {:{fill}}   - {}".format(item[0], item[1], fill=length))


def write_pot_list():
    print('List of available potential functions:\n')
    items = []
    for name, obj in inspect.getmembers(functions):
        if inspect.isclass(obj):
            myname = obj.__name__
            if myname != 'potfit_function' and not myname.startswith('stiweb') and not myname.startswith('tersoff'):
                items.append(obj.__name__)
    print(textwrap.fill(', '.join(items)))


def create_pot_table(args):
    table = []
    for pot in args.functions.split(','):
        items = pot.split('*')
        if len(items) == 1:
            count = 1
            name = pot
        elif len(items) == 2:
            try:
                count = int(items[0])
            except:
                sys.stderr.write('Could not read your potential list: >> {} << unknown.\n'.format(pot))
                sys.stderr.write('Please use the following syntax only: pot or 3*pot\n')
                sys.exit(1)
            name = items[1]
        else:
            sys.stderr.write('\nCould not read your potential list: >> {} << unknown.\n'.format(pot))
            sys.stderr.write('Please use the following syntax only: pot or 3*pot\n')
            sys.exit(1)
        do_smooth = 0
        if name.endswith('_sc'):
            do_smooth = 1
            name = name[:-3]
            if args.global_cutoff:
                do_smooth = 2
        for j in range(count):
           try:
                table.append(getattr(functions, name)())
                table[-1].set_properties(args.cutoff, do_smooth, args.random)
           except:
                sys.stderr.write("Could not find the function type >> %s <<\n\n" % name)
                write_pot_list()
                sys.exit(1)
    return table


def generate_stiweb_potential(args, interaction, output):
    table = []
    cols = interaction[2](args.ntypes)
    for i in range(paircol(args.ntypes)):
        table.append(functions.stiweb_2(args.cutoff))
        table[-1].set_properties(args.cutoff,0,args.random)
    for i in range(paircol(args.ntypes)):
        table.append(functions.stiweb_3(args.cutoff))
        table[-1].set_properties(args.cutoff,0,args.random)
    write_header(args, cols, output)
    for i in table:
        output.write('\n')
        i.write_pot_table(output)
    output.write('\n')
    output.write('type lambda\n')
    output.write('cutoff {}\n'.format(args.cutoff))
    for i in range(args.ntypes):
        for j in range(args.ntypes):
            for k in range(j,args.ntypes):
                output.write("lambda_{}{}{}\t2\t0\t3".format(i,j,k))


def generate_tersoff_potential(args, interaction, output):
    table = []
    cols = interaction[2](args.ntypes)
    for i in range(paircol(args.ntypes)):
        table.append(functions.tersoff_pot())
        table[-1].set_properties(args.cutoff, 0, args.random)
    for i in range(paircol(args.ntypes) - args.ntypes):
        table.append(functions.tersoff_mix())
        table[-1].set_properties(args.cutoff, 0, args.random)
    write_header(args, cols, output)
    for i in table:
        output.write('\n')
        i.write_pot_table(output)


def generate_regular_potential(args, interaction, output):
    pot_table = create_pot_table(args)
    cols = interaction[2](args.ntypes)
    if len(pot_table) != cols:
        sys.stderr.write('Wrong number of potential functions given!\n')
        sys.stderr.write('A {} potential for {} atom types requires {} functions, you specified {}.\n'.format(args.interaction, args.ntypes, cols, len(pot_table)))
        sys.exit(1)
    write_header(args, cols, output)
    if args.cp and args.type == 'pair':
        write_cp(n, output)
    if args.global_cutoff:
        write_global_cutoff(args.random, output)
    for i in pot_table:
        output.write('\n')
        i.write_pot_table(output)


def write_header(args, cols, output):
    output.write('#F 0 {}\n'.format(cols))
    output.write('#T {}\n'.format(args.interaction.upper()))
    if args.elements:
        output.write('#C {}\n'.format(' '.join([s for s in args.elements.split(',')])))
    output.write('#I' + cols * '{}'.format(' 0') + '\n')
    output.write('#E\n')


def write_cp(n, output):
    output.write('\n')
    for i in range(n):
        output.write('cp_{} -1 10 0'.format(i))


def write_global_cutoff(rand, output):
    output.write('\n')
    output.write('global 1\n')
    if rand:
        output.write('h {:.2f} {:.2f} {:.2f}\n'.format(0.5+1.5*random.random(),0.5,2))
    else:
        output.write('h 1 0.5 2\n')


def main():
    parser = argparse.ArgumentParser(formatter_class=argparse.RawDescriptionHelpFormatter,
      description = 'Create an analytic potential file for potift.',
      epilog= 'To specify multiple potentials you can use the following syntax:\n\n'
      '  makeapot -n 3 -i eam -f 6*eopp,3*csw,3*bjs\n\n'
      'which uses 6 eopp potentials, 3 csw and 3 bjs in this order.')

    parser.add_argument('-n', type=int, default=1, dest='ntypes', help='number of atoms types, runs from 0 to N-1')
    parser.add_argument('-c', type=float, default=6.0, dest='cutoff', help='cutoff radius (default %(default)s)')
    parser.add_argument('-g', action='store_true', dest='global_cutoff', help='use a global cutoff parameter for all potentials')
    parser.add_argument('-r', action='store_true', dest='random', help='randomize the values for the potential parameters')
    parser.add_argument('-i', type=str, dest='interaction', metavar='INTERACTION', choices=[x[0] for x in sorted(INTERACTIONS)] ,
                        help='supported interaction types are: {}'.format(', '.join([x[0] for x in sorted(INTERACTIONS)])), required=True)
    parser.add_argument('-l', '--list', action='store_true', help='list options which are available')
    parser.add_argument('--cp', action='store_true', help='enable chemical potentials (only for pair)')
    parser.add_argument('-f', type=str, dest='functions', help='comma separated list of potential functions, either name or i*name, where i=1,2,3,...')
    parser.add_argument('-o', dest='outfile', type=argparse.FileType('w'), help='write output to this file instead of stdout')
    parser.add_argument('-e', type=str, dest='elements', help='comma separated list of elements for #C header line')
    args = parser.parse_args()

    if args.list:
        write_list()
        sys.exit(0)

    # Check for sane arguments
    if args.ntypes < 1:
        sys.stderr.write("The number of atom types cannot be less than 1\n")
        sys.exit(1)

    if args.cutoff < 0:
        sys.stderr.write("The cutoff distance needs to be positive!\n")
        sys.exit(1)

    # Check interaction type
    for item in INTERACTIONS:
        if item[0] == args.interaction:
            interaction = item
            break

    if not args.outfile:
        args.outfile = sys.stdout

    if args.interaction == 'stiweb':
        generate_stiweb_potential(args, interaction, args.outfile)
    elif args.interaction == 'tersoff':
        generate_tersoff_potential(args, interaction, args.outfile)
    else:
        if not args.functions:
            sys.stderr.write('Error: Interaction {} requires list of potential functions (-f argument)\n'.format(args.interaction))
            sys.exit(1)
        generate_regular_potential(args, interaction, args.outfile)
    args.outfile.close()


if __name__ == "__main__":
    main()
